#include <gtest/gtest.h>
#include <chrono>
#include <random>
#include <vector>
#include <iostream>
#include <iomanip>
#include <algorithm>

#include "saxpy.h"
#include "gemv.h"
#include "pack.h"

class PerformanceTest : public ::testing::Test {
protected:
    void SetUp() override {
        // Fixed seed for reproducible results
        gen.seed(12345);
    }

    std::mt19937 gen;
    std::uniform_real_distribution<float> float_dist{-1000.0f, 1000.0f};
    std::uniform_int_distribution<uint8_t> uint8_dist{0, 15}; // 4-bit values
    
    // Helper function to measure execution time
    template<typename Func>
    double measure_time_ms(Func&& func, int iterations = 1) {
        auto start = std::chrono::high_resolution_clock::now();
        for (int i = 0; i < iterations; i++) {
            func();
        }
        auto end = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::nanoseconds>(end - start);
        return duration.count() / 1e6 / iterations; // Convert to milliseconds per iteration
    }
    
    // Helper function to calculate speedup
    double calculate_speedup(double scalar_time, double sve_time) {
        if (sve_time == 0.0) return 0.0;
        return scalar_time / sve_time;
    }
    
    // Helper function to print performance results
    void print_performance_result(const std::string& test_name, 
                                 double scalar_time, double sve_time, 
                                 size_t data_size = 0) {
        double speedup = calculate_speedup(scalar_time, sve_time);
        std::cout << std::fixed << std::setprecision(3);
        std::cout << "=== " << test_name << " ===" << std::endl;
        std::cout << "Scalar time:  " << scalar_time << " ms" << std::endl;
        std::cout << "SVE time:     " << sve_time << " ms" << std::endl;
        std::cout << "Speedup:      " << speedup << "x" << std::endl;
        if (data_size > 0) {
            std::cout << "Data size:    " << data_size << " elements" << std::endl;
            std::cout << "Scalar throughput: " << (data_size / (scalar_time / 1000.0)) / 1e6 << " M elements/sec" << std::endl;
            std::cout << "SVE throughput:    " << (data_size / (sve_time / 1000.0)) / 1e6 << " M elements/sec" << std::endl;
        }
        std::cout << std::endl;
    }
};

// SAXPY Performance Tests
TEST_F(PerformanceTest, SaxpyPerformanceSmallArrays) {
    const int iterations = 10000;
    const size_t n = 1000;
    const float a = 2.5f;
    
    std::vector<float> x(n), y_scalar(n), y_sve(n);
    
    // Generate random data
    for (size_t i = 0; i < n; i++) {
        x[i] = float_dist(gen);
        y_scalar[i] = y_sve[i] = float_dist(gen);
    }
    
    // Measure scalar performance
    auto scalar_backup = y_scalar;
    double scalar_time = measure_time_ms([&]() {
        y_scalar = scalar_backup;
        saxpy(x.data(), y_scalar.data(), a, n);
    }, iterations);
    
    // Measure SVE performance
    auto sve_backup = y_sve;
    double sve_time = measure_time_ms([&]() {
        y_sve = sve_backup;
        saxpy_sve(x.data(), y_sve.data(), a, n);
    }, iterations);
    
    print_performance_result("SAXPY Small Arrays", scalar_time, sve_time, n);
    
    // Verify correctness
    EXPECT_TRUE(std::equal(y_scalar.begin(), y_scalar.end(), y_sve.begin(),
                          [](float a, float b) { return std::abs(a - b) < 1e-5f; }));
}

TEST_F(PerformanceTest, SaxpyPerformanceLargeArrays) {
    const int iterations = 100;
    const size_t n = 1000000; // 1M elements
    const float a = 3.14f;
    
    std::vector<float> x(n), y_scalar(n), y_sve(n);
    
    // Generate random data
    for (size_t i = 0; i < n; i++) {
        x[i] = float_dist(gen);
        y_scalar[i] = y_sve[i] = float_dist(gen);
    }
    
    // Measure scalar performance
    auto scalar_backup = y_scalar;
    double scalar_time = measure_time_ms([&]() {
        y_scalar = scalar_backup;
        saxpy(x.data(), y_scalar.data(), a, n);
    }, iterations);
    
    // Measure SVE performance
    auto sve_backup = y_sve;
    double sve_time = measure_time_ms([&]() {
        y_sve = sve_backup;
        saxpy_sve(x.data(), y_sve.data(), a, n);
    }, iterations);
    
    print_performance_result("SAXPY Large Arrays", scalar_time, sve_time, n);
    
    // Verify correctness on a subset
    bool correct = true;
    for (size_t i = 0; i < std::min(n, size_t(1000)); i++) {
        if (std::abs(y_scalar[i] - y_sve[i]) > 1e-5f) {
            correct = false;
            break;
        }
    }
    EXPECT_TRUE(correct);
}

TEST_F(PerformanceTest, SaxpyPerformanceVaryingSizes) {
    const int iterations = 1000;
    std::vector<size_t> sizes = {100, 500, 1000, 5000, 10000, 50000, 100000};
    
    std::cout << "=== SAXPY Performance vs Array Size ===" << std::endl;
    std::cout << std::setw(10) << "Size" << std::setw(15) << "Scalar (ms)" 
              << std::setw(15) << "SVE (ms)" << std::setw(15) << "Speedup" << std::endl;
    std::cout << std::string(55, '-') << std::endl;
    
    for (size_t n : sizes) {
        const float a = 2.0f;
        std::vector<float> x(n), y_scalar(n), y_sve(n);
        
        for (size_t i = 0; i < n; i++) {
            x[i] = float_dist(gen);
            y_scalar[i] = y_sve[i] = float_dist(gen);
        }
        
        auto scalar_backup = y_scalar;
        double scalar_time = measure_time_ms([&]() {
            y_scalar = scalar_backup;
            saxpy(x.data(), y_scalar.data(), a, n);
        }, iterations);
        
        auto sve_backup = y_sve;
        double sve_time = measure_time_ms([&]() {
            y_sve = sve_backup;
            saxpy_sve(x.data(), y_sve.data(), a, n);
        }, iterations);
        
        double speedup = calculate_speedup(scalar_time, sve_time);
        
        std::cout << std::fixed << std::setprecision(3);
        std::cout << std::setw(10) << n << std::setw(15) << scalar_time 
                  << std::setw(15) << sve_time << std::setw(15) << speedup << std::endl;
    }
    std::cout << std::endl;
}

// GEMV Performance Tests
TEST_F(PerformanceTest, GemvPerformanceSquareMatrices) {
    const int iterations = 100;
    const size_t n = 500; // 500x500 matrix
    
    std::vector<float> A(n * n), x(n), y_scalar(n), y_sve(n);
    
    // Generate random data
    for (size_t i = 0; i < n * n; i++) {
        A[i] = float_dist(gen);
    }
    for (size_t i = 0; i < n; i++) {
        x[i] = float_dist(gen);
    }
    
    // Measure scalar performance
    double scalar_time = measure_time_ms([&]() {
        gemv(A.data(), x.data(), y_scalar.data(), n, n);
    }, iterations);
    
    // Measure SVE performance
    double sve_time = measure_time_ms([&]() {
        gemv_sve(A.data(), x.data(), y_sve.data(), n, n);
    }, iterations);
    
    print_performance_result("GEMV Square Matrix", scalar_time, sve_time, n * n);
}

TEST_F(PerformanceTest, GemvPerformanceRectangularMatrices) {
    const int iterations = 50;
    const size_t m = 1000, n = 500; // 1000x500 matrix
    
    std::vector<float> A(m * n), x(n), y_scalar(m), y_sve(m);
    
    // Generate random data
    for (size_t i = 0; i < m * n; i++) {
        A[i] = float_dist(gen);
    }
    for (size_t i = 0; i < n; i++) {
        x[i] = float_dist(gen);
    }
    
    // Measure scalar performance
    double scalar_time = measure_time_ms([&]() {
        gemv(A.data(), x.data(), y_scalar.data(), m, n);
    }, iterations);
    
    // Measure SVE performance
    double sve_time = measure_time_ms([&]() {
        gemv_sve(A.data(), x.data(), y_sve.data(), m, n);
    }, iterations);
    
    print_performance_result("GEMV Rectangular Matrix", scalar_time, sve_time, m * n);
}

TEST_F(PerformanceTest, GemvPerformanceVaryingSizes) {
    const int iterations = 100;
    std::vector<size_t> sizes = {50, 100, 200, 300, 400, 500};
    
    std::cout << "=== GEMV Performance vs Matrix Size ===" << std::endl;
    std::cout << std::setw(10) << "Size" << std::setw(15) << "Scalar (ms)" 
              << std::setw(15) << "SVE (ms)" << std::setw(15) << "Speedup" << std::endl;
    std::cout << std::string(55, '-') << std::endl;
    
    for (size_t n : sizes) {
        std::vector<float> A(n * n), x(n), y_scalar(n), y_sve(n);
        
        for (size_t i = 0; i < n * n; i++) {
            A[i] = float_dist(gen);
        }
        for (size_t i = 0; i < n; i++) {
            x[i] = float_dist(gen);
        }
        
        double scalar_time = measure_time_ms([&]() {
            gemv(A.data(), x.data(), y_scalar.data(), n, n);
        }, iterations);
        
        double sve_time = measure_time_ms([&]() {
            gemv_sve(A.data(), x.data(), y_sve.data(), n, n);
        }, iterations);
        
        double speedup = calculate_speedup(scalar_time, sve_time);
        
        std::cout << std::fixed << std::setprecision(3);
        std::cout << std::setw(10) << n << "x" << n << std::setw(12) << scalar_time 
                  << std::setw(15) << sve_time << std::setw(15) << speedup << std::endl;
    }
    std::cout << std::endl;
}

// Pack/Unpack Performance Tests
TEST_F(PerformanceTest, PackUnpackPerformance) {
    const int iterations = 10000;
    const size_t n = 100000; // 100K elements
    
    std::vector<uint8_t> src(n), dst_scalar(n/2), dst_sve(n/2);
    std::vector<uint8_t> unpacked_scalar(n), unpacked_sve(n);
    
    // Generate random 4-bit data
    for (size_t i = 0; i < n; i++) {
        src[i] = uint8_dist(gen);
    }
    
    // Perform packing

    pack(dst_scalar.data(), src.data(), n);
    pack(dst_sve.data(), src.data(), n);
    
    // Measure unpack performance
    double unpack_scalar_time = measure_time_ms([&]() {
        unpack(unpacked_scalar.data(), dst_scalar.data(), n);
    }, iterations);
    
    double unpack_sve_time = measure_time_ms([&]() {
        unpack_sve(unpacked_sve.data(), dst_scalar.data(), n);
    }, iterations);
    
    print_performance_result("Unpack Operation", unpack_scalar_time, unpack_sve_time, n);
    
    // Verify correctness for pack (if SVE implementation exists)
    bool pack_sve_implemented = false;
    for (size_t i = 0; i < n/2; i++) {
        if (dst_sve[i] != 0) {
            pack_sve_implemented = true;
            break;
        }
    }
    
    if (pack_sve_implemented) {
        EXPECT_TRUE(std::equal(dst_scalar.begin(), dst_scalar.end(), dst_sve.begin()));
    }
    
    // Verify correctness for unpack (if SVE implementation exists)
    bool unpack_sve_implemented = false;
    for (size_t i = 0; i < n; i++) {
        if (unpacked_sve[i] != 0) {
            unpack_sve_implemented = true;
            break;
        }
    }
    
    if (unpack_sve_implemented) {
        EXPECT_TRUE(std::equal(unpacked_scalar.begin(), unpacked_scalar.end(), unpacked_sve.begin()));
    }
}

TEST_F(PerformanceTest, PackUnpackVaryingSizes) {
    const int iterations = 1000;
    std::vector<size_t> sizes = {1000, 5000, 10000, 50000, 100000, 500000};
    
    std::cout << "=== Pack Performance vs Data Size ===" << std::endl;
    std::cout << std::setw(10) << "Size" << std::setw(15) << "Scalar (ms)" 
              << std::setw(15) << "SVE (ms)" << std::setw(15) << "Speedup" << std::endl;
    std::cout << std::string(55, '-') << std::endl;
    
    for (size_t n : sizes) {
        std::vector<uint8_t> src(n), dst_scalar(n/2), dst_sve(n/2),
            unpacked_scalar(n), unpacked_sve(n);
        
        for (size_t i = 0; i < n; i++) {
            src[i] = uint8_dist(gen);
        }
        
        // Perform packing
        pack(dst_scalar.data(), src.data(), n);
        pack(dst_sve.data(), src.data(), n);

        double unpack_scalar_time = measure_time_ms([&]() {
            unpack(unpacked_scalar.data(), dst_scalar.data(), n);
        }, iterations);

        double unpack_sve_time = measure_time_ms([&]() {
            unpack_sve(unpacked_sve.data(), dst_scalar.data(), n);
        }, iterations);
        
        double speedup = calculate_speedup(unpack_scalar_time, unpack_sve_time);
        
        std::cout << std::fixed << std::setprecision(3);
        std::cout << std::setw(10) << n << std::setw(15) << unpack_scalar_time 
                  << std::setw(15) << unpack_sve_time << std::setw(15) << calculate_speedup(unpack_scalar_time, unpack_sve_time) << std::endl;
    }
    std::cout << std::endl;
}

// Combined Performance Test
TEST_F(PerformanceTest, OverallPerformanceSummary) {
    std::cout << "=== Overall Performance Summary ===" << std::endl;
    
    // SAXPY summary test
    const size_t saxpy_n = 100000;
    const float a = 2.0f;
    std::vector<float> x(saxpy_n), y_scalar(saxpy_n), y_sve(saxpy_n);
    
    for (size_t i = 0; i < saxpy_n; i++) {
        x[i] = float_dist(gen);
        y_scalar[i] = y_sve[i] = float_dist(gen);
    }
    
    auto saxpy_scalar_backup = y_scalar;
    double saxpy_scalar_time = measure_time_ms([&]() {
        y_scalar = saxpy_scalar_backup;
        saxpy(x.data(), y_scalar.data(), a, saxpy_n);
    }, 1000);
    
    auto saxpy_sve_backup = y_sve;
    double saxpy_sve_time = measure_time_ms([&]() {
        y_sve = saxpy_sve_backup;
        saxpy_sve(x.data(), y_sve.data(), a, saxpy_n);
    }, 1000);
    
    // GEMV summary test
    const size_t gemv_n = 300;
    std::vector<float> A(gemv_n * gemv_n), gemv_x(gemv_n), gemv_y_scalar(gemv_n), gemv_y_sve(gemv_n);
    
    for (size_t i = 0; i < gemv_n * gemv_n; i++) {
        A[i] = float_dist(gen);
    }
    for (size_t i = 0; i < gemv_n; i++) {
        gemv_x[i] = float_dist(gen);
    }
    
    double gemv_scalar_time = measure_time_ms([&]() {
        gemv(A.data(), gemv_x.data(), gemv_y_scalar.data(), gemv_n, gemv_n);
    }, 100);
    
    double gemv_sve_time = measure_time_ms([&]() {
        gemv_sve(A.data(), gemv_x.data(), gemv_y_sve.data(), gemv_n, gemv_n);
    }, 100);
    
    // Pack summary test
    const size_t pack_n = 100000;
    std::vector<uint8_t> pack_src(pack_n), pack_dst_scalar(pack_n/2), pack_dst_sve(pack_n/2);
    
    for (size_t i = 0; i < pack_n; i++) {
        pack_src[i] = uint8_dist(gen);
    }
    
    // Perform packing
    pack(pack_dst_scalar.data(), pack_src.data(), pack_n);
    pack(pack_dst_sve.data(), pack_src.data(), pack_n);

    double unpack_scalar_time = measure_time_ms([&]() {
        unpack(pack_src.data(), pack_dst_scalar.data(), pack_n);
    }, 1000);
    double unpack_sve_time = measure_time_ms([&]() {
        unpack_sve(pack_src.data(), pack_dst_sve.data(), pack_n);
    }, 1000);
    
    // Print summary
    std::cout << std::fixed << std::setprecision(2);
    std::cout << "Function    | Scalar Time | SVE Time | Speedup" << std::endl;
    std::cout << std::string(45, '-') << std::endl;
    std::cout << "SAXPY       | " << std::setw(9) << saxpy_scalar_time << " ms | " 
              << std::setw(6) << saxpy_sve_time << " ms | " 
              << std::setw(5) << calculate_speedup(saxpy_scalar_time, saxpy_sve_time) << "x" << std::endl;
    std::cout << "GEMV        | " << std::setw(9) << gemv_scalar_time << " ms | " 
              << std::setw(6) << gemv_sve_time << " ms | " 
              << std::setw(5) << calculate_speedup(gemv_scalar_time, gemv_sve_time) << "x" << std::endl;
    std::cout << "Unpack        | " << std::setw(9) << unpack_scalar_time << " ms | " 
              << std::setw(6) << unpack_sve_time << " ms | " 
              << std::setw(5) << calculate_speedup(unpack_scalar_time, unpack_sve_time) << "x" << std::endl;
    std::cout << std::endl;
    
    // Performance expectations
    double saxpy_speedup = calculate_speedup(saxpy_scalar_time, saxpy_sve_time);
    double gemv_speedup = calculate_speedup(gemv_scalar_time, gemv_sve_time);
    
    EXPECT_GT(saxpy_speedup, 1.0) << "SAXPY SVE should be faster than scalar";
    EXPECT_GT(gemv_speedup, 1.0) << "GEMV SVE should be faster than scalar";
    
    std::cout << "Performance test completed successfully!" << std::endl;
} 